import sys
from time import sleep
import serial
from serial.tools import list_ports

class SerialManager:
    def __init__(self, baudrate=9600):
        self.serial = serial.Serial()
        self.serial.baudrate = baudrate
        self.serial.timeout = 0.1
        self.connected = False
    
    def get_available_ports(self):
        ports = list_ports.comports()
        available_ports = [port.device for port in ports]
        return available_ports
    
    def connect(self, port):
        self.serial.port = port
        
        self.serial.open()
        #! Important to wait for the serial to open
        sleep(1)
        if self.serial.is_open:
            self.connected = True
            return True
        else:
            self.serial.close()
            self.connected = False
            return False
    
    def disconnect(self):
        self.serial.close()
        self.connected = False
    
    def send_data(self, message):
        """send data to serial port"""
        if not self.connected:
            print("Error: Not connected to any serial port.")
            return False
        
        try:
            self.serial.write(message.encode())
            print(f"Data sent: {message}")
            return True
        except Exception as e:
            print(f"Error while sending data: {e}")
            return False
    
    def receive_data(self):
        """receive data from serial port"""
        if not self.connected:
            print("Error: Not connected to any serial port.")
            return None
        
        try:
            data = self.serial.readline()  # 读取一行数据
            if data:
                decoded_data = data.decode().strip()  # 解码并去除多余空白字符
                print(f"Data received: {decoded_data}")
                return decoded_data
            else:
                print("No data received.")
                return None
        except Exception as e:
            print(f"Error while receiving data: {e}")
            return None
            
    # =====================================================================
    # =========================== Motor Control ===========================
    # =====================================================================
    def move_motor(self, direction=1, steps=1000, speed=1000, motor_number=0):
        if self.connected:
            if direction:
                message = f"M{motor_number}:{steps}:{speed}\n"
            else:
                message = f"M{motor_number}:{steps}:-{speed}\n"
            self.send_data(message)
        else:
            print("Serial not connected")
            
    def continue_move(self, direction, speed, motor_number=0):
        if speed == 'quick':
            motor_speed = 1000
        else:
            motor_speed = 100
        if self.connected:
            if direction:
                self.move_motor(1, -1, motor_speed, motor_number)
            else:
                self.move_motor(0, -1, motor_speed, motor_number)
        else:
            print("Serial not connected")

    # =====================================================================
    # =========================== Pump Control ===========================
    # =====================================================================
    def pump_enable(self):
        self.send_data("#W0,1\n")

    def pump_disable(self):
        self.send_data("#W0,0\n")
        
    def pump_set_powerLimit(self, power=1000):
        # ppower: 0-1400 mW
        self.send_data(f"#W1,{power}\n")

    def pump_streaming_mode_enable(self):
        self.send_data("#W2,1\n")
    
    def pump_streaming_mode_disable(self):
        self.send_data("#W2,0\n")
    
    def pump_streaming_data(self):
        # #S<PUMP_ENABLED>,<VOLTAGE>,<CURRENT>,<FREQUENCY>,0,<DIGITAL_PRESSURE>,<ANA3>,0,<CHK>\n
        data = self.receive_data()
        if data is None:
            return None
        data = data.strip().strip("#S")
        
        fields = data.split(',')

        try:
            pump_enabled = int(fields[0])  # 泵启用状态
            voltage = float(fields[1])     # 电压
            current = float(fields[2])     # 电流
            frequency = float(fields[3])   # 频率
            digital_pressure = float(fields[5])  # 数字压力
            ana3 = float(fields[6])        # 模拟信号值 ANA3
            chk = fields[8]                # 校验码
            
            return [pump_enabled, voltage, current, frequency, digital_pressure, ana3, chk]
        except (IndexError, ValueError) as e:
            print(f"Error parsing data: {e}")
            return None
    
    def pump_measure(self, name):
        if name == 'drive voltage':
            # 0-60 V
            self.send_data("#R3\n")
            data = self.receive_data()
            return data.split(",")[1]
        elif name == 'drive current':
            # 0-150 mA
            self.send_data("#R4\n")
            data = self.receive_data()
            return data.split(",")[1]
        elif name == 'drive power':
            # 0-1400 mW
            self.send_data("#R5\n")
            data = self.receive_data()
            return data.split(",")[1]
        elif name == 'drive frequency':
            # 20000-23000 Hz
            self.send_data("#R6\n")
            data = self.receive_data()
            return data.split(",")[1]
    
    def pump_set_mode(self, mode):
        if mode == 'Manual':
            self.send_data("#W10,0\n")
        elif mode == 'PID':
            self.send_data("#W10,1\n")
        elif mode == 'Bang Bang':
            self.send_data("#W10,2\n")
        else:
            print("Invalid mode")
    
    def pump_set_manual_mode(self, mode):
        # 0 = Set val (register 23)
        if mode == 0:
            self.send_data("#W11,0\n")
        # 1 = Analog 1† (dial on evak kit)
        elif mode == 1:
            self.send_data("#W11,1\n")
        # 2 = Analog 2† (pr. sense on eval kit)
        elif mode == 2:
            self.send_data("#W11,2\n")
        # 3 = Analog 3 (main analog input)
        elif mode == 3:
            self.send_data("#W11,3\n")
        else:
            print("Invalid manual mode")
    
    def pump_set_manual_power(self, power):
        # 0-1400 mW
        power = int(power)
        self.send_data(f"#W23,{power}\n")
    
    
if __name__ == "__main__":
    SerialManager = SerialManager(baudrate=115200)
    # 可用串口列表
    print("可用的串口: ", SerialManager.get_available_ports())
    
    # 连接串口
    SerialManager.connect("COM5")
    print("连接状态: ", SerialManager.connected)
    
    # 发送数据
    SerialManager.send_data("#R1\n")
    # 接收数据
    SerialManager.receive_data()

    # print(SerialManager.connected)

    # test
    
    # measure
    data = SerialManager.pump_measure("drive voltage")
    print(data)
    
    # # streaming mode
    SerialManager.pump_streaming_mode_enable()
    fields = SerialManager.pump_streaming_data()
    print(fields)
    
    # for i in range(10):
    #     SerialManager.receive_data()